package com.mkx.core.lock;

import com.mkx.properties.RedisLockProperties;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.util.Assert;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

/**
 * Redis分布式批量锁实现
 * 支持同时获取和释放多个锁，提供原子性的批量操作
 */
public class BatchRedisLock implements AutoCloseable {

    private final RedisTemplate<String, Object> redisTemplate;
    private final RedisLockProperties properties;
    
    // 已获取的锁集合，用于后续释放
    private final Map<String, RedisDistributedLock> acquiredLocks = new ConcurrentHashMap<>();
    
    /**
     * 构造函数
     */
    public BatchRedisLock(RedisTemplate<String, Object> redisTemplate,
                         RedisLockProperties properties) {
        this.redisTemplate = redisTemplate;
        this.properties = properties;
    }
    
    /**
     * 构造函数，支持初始化时传入锁键
     */
    public BatchRedisLock(RedisTemplate<String, Object> redisTemplate,
                         RedisLockProperties properties,
                         String... lockKeys) {
        this(redisTemplate, properties);
    }
    
    /**
     * 同时尝试获取多个锁
     * @param lockKeys 锁键集合
     * @param timeout 超时时间
     * @param unit 时间单位
     * @return 是否成功获取所有锁
     * @throws InterruptedException 中断异常
     */
    public boolean tryLock(List<String> lockKeys, long timeout, TimeUnit unit) throws InterruptedException {
        Assert.notEmpty(lockKeys, "Lock keys cannot be empty");
        
        // 去重并排序，确保获取锁的顺序一致，避免死锁
        List<String> sortedUniqueKeys = lockKeys.stream()
                .distinct()
                .sorted()
                .collect(Collectors.toList());
        
        long startTime = System.currentTimeMillis();
        long timeoutMillis = unit.toMillis(timeout);
        
        // 用于跟踪获取失败的锁，以便在需要时释放已获取的锁
        Set<String> failedKeys = new HashSet<>();
        
        try {
            for (String lockKey : sortedUniqueKeys) {
                // 检查是否还有足够的时间继续尝试获取锁
                long elapsedTime = System.currentTimeMillis() - startTime;
                if (elapsedTime >= timeoutMillis) {
                    failedKeys.add(lockKey);
                    break;
                }
                
                // 计算剩余超时时间
                long remainingTimeout = timeoutMillis - elapsedTime;
                
                // 创建并尝试获取锁
                RedisDistributedLock lock = new RedisDistributedLock(
                        redisTemplate, lockKey, properties.getExpireTime(), 
                        properties.getRetryInterval(), properties.isAdaptiveRetryEnabled());
                
                if (lock.tryLock(remainingTimeout, TimeUnit.MILLISECONDS)) {
                    acquiredLocks.put(lockKey, lock);
                } else {
                    failedKeys.add(lockKey);
                    break;
                }
            }
            
            // 如果有失败的锁，释放已获取的所有锁
            if (!failedKeys.isEmpty()) {
                releaseAll();
                return false;
            }
            
            return true;
        } catch (Exception e) {
            // 发生异常时释放所有已获取的锁
            releaseAll();
            throw e;
        }
    }
    
    /**
     * 同时释放所有已获取的锁
     */
    public void releaseAll() {
        List<Exception> exceptions = new ArrayList<>();
        
        // 按照与获取相反的顺序释放锁，避免死锁
        acquiredLocks.keySet().stream()
                .sorted(Comparator.reverseOrder())
                .forEach(lockKey -> {
                    RedisDistributedLock lock = acquiredLocks.get(lockKey);
                    if (lock != null) {
                        try {
                            lock.unlock();
                        } catch (Exception e) {
                            exceptions.add(e);
                        }
                    }
                });
        
        // 清空已获取的锁集合
        acquiredLocks.clear();
        
        // 如果有异常，抛出第一个异常
        if (!exceptions.isEmpty()) {
            throw new RuntimeException("Error releasing locks", exceptions.get(0));
        }
    }
    
    /**
     * 检查是否持有指定的锁
     * @param lockKey 锁键
     * @return 是否持有
     */
    public boolean holdsLock(String lockKey) {
        return acquiredLocks.containsKey(lockKey);
    }
    
    /**
     * 获取当前持有的锁数量
     */
    public int getLockedCount() {
        return acquiredLocks.size();
    }
    
    /**
     * 获取当前持有的所有锁键
     */
    public Set<String> getLockedKeys() {
        return new HashSet<>(acquiredLocks.keySet());
    }
    
    /**
     * 自动关闭接口实现，支持try-with-resources语法
     */
    public void close() {
        releaseAll();
    }
}